package com.forest.util;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.HexUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.*;
import cn.hutool.crypto.symmetric.AES;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import java.nio.ByteBuffer;
import java.security.Provider;
import java.security.SecureRandom;
import java.security.Security;

/**
 * AES加解密工具类(GCM模式).
 *
 * @author Administrator
 */
public class AesUtil {

    private static final Logger LOG = LoggerFactory.getLogger(AesUtil.class);
    /**
     * GCM模式下的认证标签长度，推荐为128位
     */
    private static final Integer GCM_AUTH_TAG_LENGTH_BYTES = 128;

    private static final String ALGORITHM = "AES";
    private static final String PROVIDER = "BC";
    private static final String ALGORITHM_MODE_PADDING = "AES/GCM/NoPadding";
    /** 混淆参数长度 */
    private static final Integer IV_LENGTH_BYTES = 12;

    /** 加密nonce长度 */
    private static final Integer NONCE_LENGTH_BYTES = 12;

    /** 密匙key长度 */
    private static final Integer KEY_LENGTH = 32;

    /** 是否使用国密提供者 */
    private static boolean useBouncyCastle;

    static {
        try {
            //默认使用国密提供者
            useBouncyCastle = EnvironmentHolder.getUseBouncyCastle();
            //初始化BouncyCastleProvider
            Provider provider = Security.getProvider(PROVIDER);
            if (provider == null) {
                Security.addProvider(new BouncyCastleProvider());
            }
        } catch (Exception e) {
            LOG.error("初始化失败: {}", e.getMessage(), e);
            if(useBouncyCastle && EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("初始化国密提供者失败");
            }
        }
    }

    /**
     * 设置是否使用国密提供者.
     *
     * @param useBouncyCastle
     */
    public static void setUseBouncyCastle(boolean useBouncyCastle) {
        AesUtil.useBouncyCastle = useBouncyCastle;
    }

    public static byte[] encrypt(String data, String keyStr) {
        return encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), keyStr, null, true);
    }

    public static byte[] encrypt(byte[] data, String keyStr, String nonceStr, boolean wrap) {
        //创建key对象
        String secretKeyStr = StrUtil.isBlank(keyStr) ? EnvironmentHolder.getSecretKey() : keyStr;
        if(StrUtil.isBlank(secretKeyStr)) {
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("请设置密钥");
            } else {
                return data;
            }
        }
        byte[] key = SecureUtil.decode(secretKeyStr);
        if(key.length != KEY_LENGTH) {
            LOG.error("密钥长度错误");
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("密钥长度错误");
            } else {
                return data;
            }
        }
        SecretKey secretKey = KeyUtil.generateKey(ALGORITHM, key);
        byte [] nonce = StrUtil.isBlank(nonceStr) ? randomBytes(NONCE_LENGTH_BYTES) : SecureUtil.decode(nonceStr);
        //创建加密参数
        GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(GCM_AUTH_TAG_LENGTH_BYTES, nonce);
        try {
            Cipher cipher;
            if(useBouncyCastle) {
                cipher = Cipher.getInstance(ALGORITHM_MODE_PADDING, PROVIDER);
            } else {
                cipher = Cipher.getInstance(ALGORITHM_MODE_PADDING);
            }
            cipher.init(Cipher.ENCRYPT_MODE, secretKey, gcmParameterSpec);
            byte[] encrypted = cipher.doFinal(data);
            return wrap ? wrap(nonce, encrypted, StrUtil.isNotBlank(nonceStr)) : encrypted;
        } catch (Exception e) {
            LOG.error("加密失败: {}", e.getMessage(), e);
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("加密失败: " + e.getMessage());
            } else {
                return data;
            }
        }
    }

    /**
     * 封装数据
     * @param nonce
     * @param encrypted
     * @return
     */
    private static byte[] wrap(byte[] nonce, byte[] encrypted, boolean fixedIv) {
        ByteBuffer buffer;
        if(fixedIv) {
            buffer = ByteBuffer.allocate(NONCE_LENGTH_BYTES + encrypted.length);
            buffer.put(nonce, 0, NONCE_LENGTH_BYTES);
            buffer.put(encrypted);
        } else {
            byte[] iv = randomBytes(IV_LENGTH_BYTES);
            buffer = ByteBuffer.allocate(NONCE_LENGTH_BYTES + IV_LENGTH_BYTES + encrypted.length);
            buffer.put(nonce, 0, NONCE_LENGTH_BYTES);
            buffer.put(encrypted);
            buffer.put(iv,0, IV_LENGTH_BYTES);
        }
        return buffer.array();
    }

    public static byte[] randomBytes(int length) {
        byte[] bytes = new byte[length];
        SecureRandom random = new SecureRandom();
        random.nextBytes(bytes);
        return bytes;
    }

    public static String encryptBase64(String data) {
        byte[] encrypt = encrypt(data, null);
        return Base64.encode(encrypt);
    }

    public static String encryptBase64(String data,String key) {
        byte[] encrypt = encrypt(data,key);
        return Base64.encode(encrypt);
    }

    public static String encryptHex(String data) {
        byte [] encrypt = encrypt(data, null);
        return HexUtil.encodeHexStr(encrypt);
    }

    public static String encryptHex(String data,String key) {
        byte [] encrypt = encrypt(data, key);
        return HexUtil.encodeHexStr(encrypt);
    }

    /**
     * 解密.
     *
     * @param data
     * @return
     */
    public static byte[] decrypt(byte[] data, String keyStr, String nonceStr, boolean wrap) {
        //创建key对象
        String secretKeyStr = StrUtil.isBlank(keyStr) ? EnvironmentHolder.getSecretKey() : keyStr;
        if(StrUtil.isBlank(secretKeyStr)) {
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("请设置密钥");
            } else {
                return data;
            }
        }
        byte[] key = SecureUtil.decode(secretKeyStr);
        if(key.length != KEY_LENGTH) {
            LOG.error("密钥长度错误");
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("密钥长度错误");
            } else {
                return data;
            }
        }
        SecretKey secretKey = KeyUtil.generateKey(ALGORITHM, key);
        //创建加密参数
        byte[] nonce = StrUtil.isBlank(nonceStr) ? unwrapNonce(data) : SecureUtil.decode(nonceStr);
        //解密数据
        byte[] encrypted = unwrapEncrypted(data, StrUtil.isNotBlank(nonceStr), wrap);
        //创建加密参数
        GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(GCM_AUTH_TAG_LENGTH_BYTES, nonce);
        try {
            Cipher cipher;
            if(useBouncyCastle) {
                cipher = Cipher.getInstance(ALGORITHM_MODE_PADDING, PROVIDER);
            } else {
                cipher = Cipher.getInstance(ALGORITHM_MODE_PADDING);
            }
            cipher.init(Cipher.DECRYPT_MODE, secretKey, gcmParameterSpec);
            byte[] decrypt = cipher.doFinal(encrypted);
            return decrypt;
        } catch (Exception e) {
            LOG.error("解密失败: {}", e.getMessage(), e);
            if(EnvironmentHolder.getThrowException()) {
                throw new RuntimeException("解密失败: " + e.getMessage());
            } else {
                return data;
            }
        }
    }

    private static byte[] unwrapEncrypted(byte[] data, boolean fixedIv, boolean wrap) {
        if(!wrap) {
            return data;
        }
        int dataLength = fixedIv ? data.length - NONCE_LENGTH_BYTES : data.length - NONCE_LENGTH_BYTES - IV_LENGTH_BYTES;
        byte[] encrypted = new byte[dataLength];
        System.arraycopy(data, NONCE_LENGTH_BYTES , encrypted, 0, encrypted.length);
        return encrypted;
    }

    private static byte[] unwrapNonce(byte[] data) {
        byte [] nonce = new byte[NONCE_LENGTH_BYTES];
        System.arraycopy(data, 0, nonce, 0, NONCE_LENGTH_BYTES);
        return nonce;
    }

    /**
     * 解密(hex/base64)
     * @param data
     * @return
     */
    public static String decrypt(String data) {
        byte[] decode = SecureUtil.decode(data);
        byte[] decrypt = decrypt(decode, null, null, true);
        return decrypt == null ? data : new String(decrypt);
    }

    public static String decrypt(String data, String keyStr) {
        byte[] decode = SecureUtil.decode(data);
        byte[] decrypt = decrypt(decode, keyStr, null, true);
        return decrypt == null ? data : new String(decrypt);
    }

    public static String decrypt(String data, String keyStr, String nonceStr) {
        byte[] decode = SecureUtil.decode(data);
        byte[] decrypt = decrypt(decode, keyStr, nonceStr, true);
        return decrypt == null ? data : new String(decrypt);
    }

    public static String decrypt(String data, String keyStr, String nonceStr, boolean wrap) {
        byte[] decode = SecureUtil.decode(data);
        byte[] decrypt = decrypt(decode, keyStr, nonceStr, wrap);
        return decrypt == null ? data : new String(decrypt);
    }

    public static String encryptBase64(String data, String key, String iv) {
        byte[] encrypt = encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), key, iv, true);
        return Base64.encode(encrypt);
    }

    public static String encryptBase64(String data, String key, String iv, boolean wrap) {
        byte[] encrypt = encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), key, iv, wrap);
        return Base64.encode(encrypt);
    }

    public static String encryptHex(String data, String key, String iv) {
        byte[] encrypt = encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), key, iv, true);
        return HexUtil.encodeHexStr(encrypt);
    }

    public static String encryptHex(String data, String key, String iv, boolean wrap) {
        byte[] encrypt = encrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8), key, iv, wrap);
        return HexUtil.encodeHexStr(encrypt);
    }
}
